import cv2
import numpy as np
import matplotlib.pyplot as plt

def calculate_phash(image_path, hash_size=8):
    """
    计算图片的pHash（感知哈希值，基于DCT）。

    Args:
        image_path (str): 图片文件的路径。
        hash_size (int: 哈希尺寸，用于DCT后的低频区域大小，常用8或16。

    Returns:
        str: 图片的pHash值（二进制字符串）。
    """
    # 1. 加载图片并转换为灰度图
    img = cv2.imread(image_path)
    if img is None:
        raise FileNotFoundError(f"无法加载图片: {image_path}")

    img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # 2. 缩小图片尺寸到hash_size * hash_size
    # 缩小是为了减少DCT的计算量，并保留主要结构信息
    img_resized = cv2.resize(img_gray, (hash_size * 8, hash_size * 8), interpolation=cv2.INTER_AREA)

    # 3. 进行DCT变换
    # DCT会将图像从空间域转换到频率域
    dct = cv2.dct(np.float32(img_resized))

    # 4. 提取左上角低频分量
    # 低频分量集中在DCT矩阵的左上角，代表图像的整体结构
    dct_low_freq = dct[0:hash_size, 0:hash_size]

    # 5. 计算低频区域的平均值
    # 作为区分高低值的阈值
    mean_val = np.mean(dct_low_freq)

    # 6. 生成哈希指纹
    # 遍历低频区域，高于平均值为1，否则为0
    phash_bits = (dct_low_freq >= mean_val).flatten()
    phash_str = "".join(['1' if b else '0' for b in phash_bits])

    return phash_str

def hamming_distance(hash1, hash2):
    """
    计算两个哈希字符串之间的汉明距离。
    Args:
        hash1 (str): 第一个哈希字符串。
        hash2 (str): 第二个哈希字符串。
    Returns:
        int: 汉明距离（不同位的数量）。
    """
    if len(hash1) != len(hash2):
        raise ValueError("哈希字符串长度不一致，无法计算汉明距离。")
    distance = sum(c1 != c2 for c1, c2 in zip(hash1, hash2))
    return distance

if __name__ == "__main__":
    # --- 替换为你的图片路径 ---
    original_image_path = './Image_deduplication/original_image.jpg'
    different_image_path = './Image_deduplication/different.jpg'
    # -------------------------

    print("--- pHash DEMO 开始 ---")

    try:
        hash_original = calculate_phash(original_image_path)
        hash_different = calculate_phash(different_image_path)

        print(f"原始图片哈希 ({original_image_path}): {hash_original}")
        print(f"不同图片哈希 ({different_image_path}): {hash_different}")
        print("-" * 30)

        # 比较哈希值
        dist_original_different = hamming_distance(hash_original, hash_different)

        print(f"原始图片 vs 不同图片 汉明距离: {dist_original_different} (期望值很高)")

        # 经验判断阈值：通常汉明距离小于5-10认为图片相似
        # 对于“不一样的图片”，我们期望汉明距离会远大于这个阈值
        threshold = 10
        print("\n--- 相似性判断 (阈值: 10) ---")
        print(f"原始图片 vs 不同图片 相似性: {'相似' if dist_original_different < threshold else '不相似'}")

        # 可视化 pHash 处理过程
        img = cv2.imread(original_image_path)
        if img is None:
            raise FileNotFoundError(f"无法加载图片: {original_image_path}")

        img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        hash_size = 8
        img_resized = cv2.resize(img_gray, (hash_size * 8, hash_size * 8), interpolation=cv2.INTER_AREA)
        dct = cv2.dct(np.float32(img_resized))
        dct_low_freq = dct[0:hash_size, 0:hash_size]
        mean_val = np.mean(dct_low_freq)
        phash_bits = (dct_low_freq >= mean_val).flatten()
        phash_img = np.array(phash_bits).reshape((hash_size, hash_size))

        fig, axes = plt.subplots(1, 3, figsize=(15, 5))
        axes[0].imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        axes[0].set_title('Original Image')
        axes[0].axis('off')

        axes[1].imshow(img_resized, cmap='gray')
        axes[1].set_title('Processed Image (before pHash)')
        axes[1].axis('off')

        axes[2].imshow(phash_img, cmap='binary')
        axes[2].set_title('Binary Image (pHash)')
        axes[2].axis('off')

        plt.show()

    except FileNotFoundError as e:
        print(f"错误: {e}. 请检查图片路径是否正确，以及文件是否存在。")
    except Exception as e:
        print(f"发生未知错误: {e}")

    print("--- pHash DEMO 结束 ---")